{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.855Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import math\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# --------------------- ECA模块定义 ---------------------\n",
    "class ECALayer(nn.Module):\n",
    "    def __init__(self, channels, gamma=2, b=1):\n",
    "        super(ECALayer, self).__init__()\n",
    "        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n",
    "        # 动态计算卷积核大小\n",
    "        kernel_size = int(abs((math.log(channels, 2) + b) / gamma))\n",
    "        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1\n",
    "        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size-1)//2, bias=False)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # 通道注意力权重计算\n",
    "        y = self.avg_pool(x)            # [B, C, 1, 1]\n",
    "        y = y.squeeze(-1).transpose(1, 2)  # [B, 1, C]\n",
    "        y = self.conv(y)                # 1D卷积捕获通道关系\n",
    "        y = self.sigmoid(y)              # [B, 1, C]\n",
    "        y = y.transpose(1, 2).unsqueeze(-1)  # [B, C, 1, 1]\n",
    "        return x * y.expand_as(x)        # 特征图加权"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# --------------------- 精简版ResNet50+ECA ---------------------\n",
    "class ECABottleneck(nn.Module):\n",
    "    expansion = 4  # 扩展系数从1改为4\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, stride=1):\n",
    "        super(ECABottleneck, self).__init__()\n",
    "        # 1x1卷积降维\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        \n",
    "        # 3x3卷积\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, \n",
    "                               stride=stride, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        \n",
    "        # 1x1卷积升维\n",
    "        self.conv3 = nn.Conv2d(out_channels, out_channels * self.expansion, \n",
    "                               kernel_size=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)\n",
    "        \n",
    "        # ECA模块放在最后一个卷积之后\n",
    "        self.eca = ECALayer(out_channels * self.expansion)\n",
    "\n",
    "        # 捷径连接\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_channels != out_channels * self.expansion:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_channels, out_channels * self.expansion, \n",
    "                          kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(out_channels * self.expansion)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.relu(self.bn1(self.conv1(x)))\n",
    "        out = torch.relu(self.bn2(self.conv2(out)))\n",
    "        out = self.bn3(self.conv3(out))\n",
    "        out = self.eca(out)  # 应用ECA注意力\n",
    "        out += self.shortcut(x)\n",
    "        return torch.relu(out)\n",
    "\n",
    "class ECA_ResNet50(nn.Module):  \n",
    "    def __init__(self, block, num_blocks, num_classes=100):  \n",
    "        super(ECA_ResNet50, self).__init__()\n",
    "        self.in_channels = 64\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)\n",
    "        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)\n",
    "        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)\n",
    "        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(512 * block.expansion, num_classes)  \n",
    "\n",
    "    def _make_layer(self, block, out_channels, num_blocks, stride):\n",
    "        strides = [stride] + [1] * (num_blocks - 1)\n",
    "        layers = []\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_channels, out_channels, stride))\n",
    "            self.in_channels = out_channels * block.expansion\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.relu(self.bn1(self.conv1(x)))\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "        x = self.avgpool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = self.fc(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# --------------------- 训练与验证代码 ---------------------\n",
    "def train(model, device, train_loader, optimizer, criterion, epoch):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    for batch_idx, (data, target) in enumerate(train_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        output = model(data)\n",
    "        loss = criterion(output, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "        _, predicted = output.max(1)\n",
    "        total += target.size(0)\n",
    "        correct += predicted.eq(target).sum().item()\n",
    "        \n",
    "        if batch_idx % 100 == 0:\n",
    "            print(f'Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '\n",
    "                  f'({100. * batch_idx / len(train_loader):.0f}%)]\\tLoss: {loss.item():.6f}')\n",
    "    \n",
    "    train_loss = total_loss / len(train_loader)\n",
    "    train_acc = 100. * correct / total\n",
    "    print(f'Train set: Average loss: {train_loss:.4f}, Accuracy: {train_acc:.2f}%')\n",
    "    return train_loss, train_acc\n",
    "\n",
    "def test(model, device, test_loader, criterion):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in test_loader:\n",
    "            data, target = data.to(device), target.to(device)\n",
    "            output = model(data)\n",
    "            test_loss += criterion(output, target).item()\n",
    "            pred = output.argmax(dim=1, keepdim=True)\n",
    "            correct += pred.eq(target.view_as(pred)).sum().item()\n",
    "    \n",
    "    test_loss /= len(test_loader)\n",
    "    acc = 100. * correct / len(test_loader.dataset)\n",
    "    print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({acc:.2f}%)')\n",
    "    return test_loss, acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# --------------------- 主函数 ---------------------\n",
    "if __name__ == '__main__':\n",
    "    # 超参数设置\n",
    "    batch_size = 128\n",
    "    epochs = 100\n",
    "    lr = 0.1\n",
    "    num_classes = 100  \n",
    "    \n",
    "    # 数据预处理\n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n",
    "    ])\n",
    "\n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),\n",
    "    ])\n",
    "\n",
    "    # 加载数据集\n",
    "    train_set = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)\n",
    "    test_set = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)\n",
    "    \n",
    "    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "    test_loader = torch.utils.data.DataLoader(test_set, batch_size=100, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "    # 初始化模型（改为ResNet50）\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    \n",
    "    # 使用ResNet50结构：[3, 4, 6, 3]个块\n",
    "    model = ECA_ResNet50(ECABottleneck, [3, 4, 6, 3], num_classes=num_classes).to(device)\n",
    "    \n",
    "    # 打印模型参数数量\n",
    "    total_params = sum(p.numel() for p in model.parameters())\n",
    "    print(f\"模型总参数: {total_params/1e6:.2f}M\")\n",
    "    \n",
    "    # 优化器和学习率调度\n",
    "    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)\n",
    "    scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[30, 60, 90], gamma=0.1)\n",
    "    criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "execution_failed": "2025-05-29T01:12:44.856Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "    # 训练循环\n",
    "    best_acc = 0.0\n",
    "    history = {'train_loss': [], 'train_acc': [], 'test_loss': [], 'test_acc': []}\n",
    "\n",
    "\n",
    "    for epoch in range(1, epochs + 1):\n",
    "        train_loss, train_acc = train(model, device, train_loader, optimizer, criterion, epoch)\n",
    "        test_loss, test_acc = test(model, device, test_loader, criterion)\n",
    "        scheduler.step()\n",
    "        \n",
    "        # 记录历史数据\n",
    "        history['train_loss'].append(train_loss)\n",
    "        history['train_acc'].append(train_acc)\n",
    "        history['test_loss'].append(test_loss)\n",
    "        history['test_acc'].append(test_acc)\n",
    "        \n",
    "        if test_acc > best_acc:\n",
    "            best_acc = test_acc\n",
    "            torch.save({\n",
    "                'epoch': epoch,\n",
    "                'model_state_dict': model.state_dict(),\n",
    "                'optimizer_state_dict': optimizer.state_dict(),\n",
    "                'accuracy': test_acc,\n",
    "            }, 'eca_resnet50_cifar100.pth')\n",
    "        \n",
    "        print(f'Epoch {epoch}/{epochs} | '\n",
    "              f'LR: {scheduler.get_last_lr()[0]:.5f} | '\n",
    "              f'Best Acc: {best_acc:.2f}%')\n",
    "    \n",
    "    print(f\"训练完成，最高测试准确率: {best_acc:.2f}%\")\n",
    "    \n",
    "    # 保存最终模型\n",
    "    torch.save(model.state_dict(), 'eca_resnet50_final_cifar100.pth')\n",
    "    # 绘制准确率曲线\n",
    "    plt.figure(figsize=(10, 4))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(history['train_acc'], label='Train Acc')\n",
    "    plt.plot(history['test_acc'], label='Test Acc')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy (%)')\n",
    "    plt.title('Accuracy Curve')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "\n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(history['train_loss'], label='Train Loss')\n",
    "    plt.plot(history['test_loss'], label='Test Loss')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('Loss Curve')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "\n",
    "    # 保存图像\n",
    "    plt.tight_layout()\n",
    "    plt.savefig('eca_resnet50_training_curves.png')  # 保存为PNG图像\n",
    "    plt.show()  "
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [],
   "dockerImageVersionId": 31040,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
